# -*- coding: utf-8 -*-
import numpy as np
#this is a class of Full Quantum Mechanical Calculation of Protein Force and Energy EE-GMFCC

class EEGMFCC():
  def setGauss(self,process='8',mem='600MW',method='hf',basis='6-31g',key=('force',\
               'nosymm')):
    self.process = process ; self.mem = mem ; self.method = method
    self.basis = basis ; self.key = key
#GJF if gauss input file.
#cube(x,y) is molecular fragment,and x is a numpy.array(n,3),x[i] is ith atom's coordinate.
#y is numpy.array(n),y[i] is ith atom's type.
#该函数用于写分子碎片的高斯输入文件。
#参数cube是一个元组，第一个元素是一个n行3列的数组，用于储存原子坐标，第二个元素是一个长度
#为n的数组，用于存储相应原子的类型。
#backgroud代表背景电荷，第一个元素是一个n行3列的数组，用于存储电荷坐标，第二个元素是长度
#为n的数组，用于存储相应电荷量。
#charge是分子碎片的带电总量。
#spin是分子碎片的自旋多重度。
#name 是高斯输入文件的名字。
#keys是一个长度为n行3列的数组，每行的第一个为关键字，第二个为用于产生与关键字有关的数据的函数
#第三个为传给函数的参数。

  def writeGJF(self,cube,charge,spin,name,*keys):
    atomCrd = cube[0];atomName = cube[1]
    gauss = open(name+'.gjf','w')  #open the gjf file
    gauss.write('%chk='+name+'.chk\n') #write chk file name
    gauss.write('%nproc='+self.process+'\n')  # assign process number
    gauss.write('%mem='+self.mem+'\n') #assign memory
    #assign method,basis set and some calculate keys.
    control = '#p '+' '.join(self.key)+' '+self.method+'/'+self.basis + ' '
    data = []#这里data是需要需要写入高斯文件中的数据，与关键字一一对应。
    for operate in keys:
      control += (operate[0]+' ')#add key to control first
      data.append(operate[1](operate[2])) #get key's data
    gauss.write(control + '\n\n')
    gauss.write('EE-GMFCC AND MOLECULAR X\n\n')
    #根据高斯输入文件的顺序，我们首先需要写入分子结构的字段，然后才能写控制关键字的数据字段。
    #write group spin and charge
    gauss.write(str(charge)+' '+str(spin)+'\n')
    #write group(name and coordinate,read gauss tutorial for detail)
    atomNu = len(cube[1])
    for i in xrange(atomNu):
      gauss.write(' {0:<8}{1:>12.7f}{2:12.7f}{3:12.7f}\n'.format(atomName[i],atomCrd[i][0],\
                   atomCrd[i][1],atomCrd[i][2]))
    gauss.write('\n')
    for m in data: #写入数据
      gauss.write(m)
    
#write backgroud charge coordinate(calculate charge's force)
#backgroud(x,y) is backgroud charge,x[i] is coordinate,y[i] is charge number.   
  def writeBackgroud(self,backgroud):
    backCrd = backgroud[0];backCharge=backgroud[1];chargeNu = len(backgroud[1])
    key = []
    for i in xrange(chargeNu):
      key.append(' {0:14.8f}{1:14.8f}{2:14.8f}{3:14.8f}\n'.format(backCrd[i][0],backCrd[i][1],\
                 backCrd[i][2],backCharge[i]))
    key.append('\n') #字段之间需要空一行。
    key = ''.join(key)
    return key
#read gauss output file and get molecular energy and force.
#name -> output file name.  
  def readGaussLog(self,name,*keys):
    log = open(name,'r')
    result = []
    for key in keys:
      result.append(key(log))
    return result

  def readEnergy(self,log):
    while 1:
      line = log.readline()
      if 'SCF Done:' in line: #location the energy line
        break
    return  float(line[23:39]) # return molecular energy

  def readChargeEnergy(self,log):
    while 1:
      line = log.readline()
      if 'Self energy of the charges' in line:
        break
    return float(line[30:50])

  def readForce(self,log):
    while 1:
      line = log.readline()
      if 'Forces (Hartrees/Bohr)' in line:
        break
    log.readline();log.readline()
    records = []
    for line in log:
      if '---' in line:
        break
      records.append(line[24:])
    records = ''.join(records)
    forces = np.fromstring(records, dtype=float, sep=' ')
    forces = forces.reshape(forces.size/3,3)
    return forces
      
  def run(self):
    self.setGauss()
    spin = 1 ; charge = 1
    name = 'aaa'
    atom = [[1,1,1],[2,2,2],[3,3,3]]
    atomName = ['H','O','H']
    cube = (atom,atomName)
    back = [[4,4,4],[5,5,5],[6,6,6]]
    backCharge = [1,1,1]
    backgroud = (back,backCharge)
    back = ('charge',self.writeBackgroud,backgroud)
    a.writeGJF(cube,charge,spin,name,back) 
    energy = 0;chargeEnergy=0;readForce=0
    result = a.readGaussLog('zn_p_001002.log',self.readChargeEnergy,self.readEnergy,self.readForce)
    print result
    
      
      
    
if __name__ == '__main__':
  a = EEGMFCC()
  a.run()
